package main

import (
	"io"
	"log"
	"net"
	"time"

	pb "rpcperftest/perf"

	"google.golang.org/grpc"
)

type testServer struct {
	pb.UnimplementedTestServer
}

func (t *testServer) DownloadRequest(downloadOption *pb.DownloadOption, stream pb.Test_DownloadRequestServer) error {
	// 测试时长
	endTime := time.Now().Add(time.Duration(downloadOption.DurationTime) * time.Second)
	var size int32 = 4096
	for {
		if err := stream.Send(&pb.DownloadData{TestData: string(make([]byte, size)), DataSize: size}); err != nil {
			return err
		}
		if time.Now().After(endTime) {
			break
		}
	}
	return nil
}

func (t *testServer) UploadRequest(stream pb.Test_UploadRequestServer) error {
	var (
		recvSize int64
		recvRate float32
	)
	startTime := time.Now()
	for {
		uploadData, err := stream.Recv()
		if err == io.EOF {
			recvSize += int64(uploadData.GetDataSize())
			costTime := time.Since(startTime).Seconds()
			recvRate = float32(float64(recvSize) / time.Since(startTime).Seconds() / 1024 / 1024)
			return stream.SendAndClose(&pb.RecvSummary{
				BytesCount: recvSize,
				CostTime:   int32(costTime),
				Rate:       recvRate,
			})
		}
		if err != nil {
			return err
		}
		recvSize += int64(uploadData.GetDataSize())
	}
}

func (t *testServer) BiTestRequest(pb.Test_BiTestRequestServer) error {
	return nil
}

func newTestServer() *testServer {
	return &testServer{}
}

func main() {
	lis, err := net.Listen("tcp", "localhost:5000")
	if err != nil {
		log.Fatalln("net.Listen err:", err)
	}
	grpcServer := grpc.NewServer()
	pb.RegisterTestServer(grpcServer, newTestServer())
	if err := grpcServer.Serve(lis); err != nil {
		log.Fatalln("grpcServer.Serve err:", err)
	}

}
